# 2023(好题, 需要记住组合数的取余方法)
"""
给定 n, m，请求出所有 n 位十进制整数中有多少个数中恰好出现了 m 个 2023。

例如 00202312023 是一个 11 位的出现了 2 个 2023 的十进制整数。由于结果可能很大，请输出答案对 998244353 取模的结果。
输入n = 5 m = 1 输出:20
"""
from leetcode import test_function


def count_m2023_in_len_n(n, m):
    """使用组合数来计算
    x x 2023 x x 2023
    那么可以将2023看作1位, 用a来代替:
    x x a x x a
    既可以用组合数来计算所有可能出现的情况, C(6,2)=15种
    aaxxxx  xaaxxx  xxaaxx  xxxaax  xxxxaa
    axaxxx  xaxaxx  xxaxax  xxxaxa
    axxaxx  xaxxax  xxaxxa
    axxxax  xaxxxa
    axxxxa

    每一位x可以使用0-9共10个数字

    """

    def combination(m, k):
        if k < 0 or k > m:
            return 0
        if k == m or k == 0:
            return 1

        k = min(k, m - k)
        res = 1
        for i in range(1, k + 1):
            res = res * (m - k + i) // i
        return res

    remain = n - 4 * m  # 出现m个2023后剩余的位数
    if remain < 0:
        return 0

    total = m + remain  # 可以将每个2023看作1位, 总计可以有total位
    c = combination(total, m)  # 所有可能出现的情况

    return c * (10 ** remain)


def count_m2023_in_len_n_opt(n, m):
    MOD = 998244353

    # 计算 a^b % mod 的快速幂
    def power_mod(a, b, mod):
        res = 1
        while b:
            if b & 1:
                res = res * a % mod
            a = a * a % mod
            b >>= 1
        return res

    # 预处理阶乘和逆元
    def precompute_factorial(n, mod):
        fact = [1] * (n + 1)
        inv_fact = [1] * (n + 1)

        for i in range(2, n + 1):
            fact[i] = fact[i - 1] * i % mod
            # 逆元计算：inv_fact[n] = (fact[n])^(mod-2) % mod
            inv_fact[i] = inv_fact[i - 1] * power_mod(i, mod - 2, mod) % mod

        return fact, inv_fact

    # 计算 C(n, k) % mod
    def comb(n, k, mod, fact, inv_fact):
        if k > n or k < 0:
            return 0
        return fact[n] * inv_fact[k] % mod * inv_fact[n - k] % mod

    remain = n - 4 * m  # 出现m个2023后剩余的位数
    if remain < 0:
        return 0

    fact, inv_fact = precompute_factorial(n, MOD)
    total = m + remain  # 可以将每个2023看作1位, 总计可以有total位
    c = comb(total, m, MOD, fact, inv_fact)  # 所有可能出现的情况
    c = c * power_mod(10, remain, MOD)
    return c % MOD


def std(n, m):
    """标准答案"""
    MOD = 998244353

    # 计算 a^b % mod 的快速幂
    def power_mod(a, b, mod):
        res = 1
        while b:
            if b & 1:
                res = res * a % mod
            a = a * a % mod
            b >>= 1
        return res

    # 预处理阶乘和逆元
    def precompute_factorial(n, mod):
        fact = [1] * (n + 1)
        inv_fact = [1] * (n + 1)

        for i in range(2, n + 1):
            fact[i] = fact[i - 1] * i % mod
            # 逆元计算：inv_fact[n] = (fact[n])^(mod-2) % mod
            inv_fact[i] = inv_fact[i - 1] * power_mod(i, mod - 2, mod) % mod

        return fact, inv_fact

    # 计算 C(n, k) % mod
    def comb(n, k, mod, fact, inv_fact):
        if k > n or k < 0:
            return 0
        return fact[n] * inv_fact[k] % mod * inv_fact[n - k] % mod

    k = n // 4
    ans = 0
    fact, inv_fact = precompute_factorial(n, MOD)
    for i in range(m, k + 1):
        ans += (-1 if (i - m) % 2 else 1) * comb(i, m, MOD, fact, inv_fact) * comb(n - 3 * i, i, MOD, fact, inv_fact,
                                                                                   ) * power_mod(10, n - 4 * i, MOD)
        ans %= MOD

    return ans


if __name__ == '__main__':
    inp = [{"n": 4, "m": 1}, {"n": 5, "m": 1}, {"n": 6, "m": 1}]
    out = [1, 20, 300]
    test_function(count_m2023_in_len_n, inp, out)
    test_function(count_m2023_in_len_n_opt, inp, out)
    test_function(std, inp, out)
